%clear; %clc; close all;
rng('default');

%%% Set Parameters
periods_within_year = 1; % benchmark model is annual

sigma_l = 0.2/sqrt(periods_within_year); % std dev of idio iid endowment shocks
lbar    =  - 0.5*sigma_l^2;
abar    = 0;
p_unemp   = 0;     % in benchmark model there is no "unemployment" so p_unemp = 0. If you set above zero, then there p_unemp chance (iid over time)
% for a period where the agent receivews only w_unempl income (i.e. unmployment insurance, not market wage)
w_unemp   = 0.05;  % wage replacement in unemployment

% households
gam = 1; % CRRA
b = 0; % borrowing constraint
deltax = 0.02/periods_within_year; % death probability
delta_eta = 0;   % How much of the information of prior generations gets inherited by offspring. Benchmark in paper is 0, meaning no information transfer across generations
delta_inh = 1;   % How much of assets gets transfered across generations. Benchmark in paper is 1, meaning full transfer of resources to new agents.
                 % you can also set delta_inh = 0, in which case no assets
                 % are inherited and new generations start with 0 assets.
                 % The estates are distributed across surviving agents
                 % using usual Blanchard (1985) style insurance markets
                 % note: values for delta_inh other than 0 and 1 are not supported

ex_signals = 0*1.3065;  % benchmark model ex_signals = 0. If it is a positive number, then the code will simulate a counter-factual where agents do not
                        % choose the optimal reasoning strategy, but rather receive a signal with exogenous noise variance equal to ex_signals each period

bta = 0.96^(1/periods_within_year);  % discount factor
btat = bta*(1-deltax*(1-delta_inh)); % delta_inh = 1 gives us the benchmark case in the manuscript. This is a model where agents care equally as much about
                                     % their offspring as they care about themselves.
                                     % when delta_inh = 0 we have a standard perpetual youth model, with an effectively lower time discount

W_cc = 1;
kappa = 0.9656;
sigma_c = sqrt(0.7658);
psi     = 0.05;
psi_tau = 0.7;          %%% this parameter is only needed in the policy experiment extension described in Appendix K
cov_fun = 1;
factor_mpc1 = 1/12;     %%%% size of shock for computing MPCs, in terms of std deviation of annual income. so the small shock is 
                        %%%% roughly one monthly std. deviation
factor_mpc2 = 1;        %%%% size of "large" shock for computing MPCs our of "large" shocks -- 1 std. deviation of annual income      

sigma_tremble = 0.183; %%%% variance of "trembles" in the control cost version of the model (Section 5.3)

[kappa, sigma_c^2, psi, b]

% firms
alpha = 0.36; % capital share
delta = 0.08/periods_within_year; % depreciation rate
Z     = 1;

util_type = 'crra';
run_fsolve   = 0;
log_c = 0;     %%%% set to 1 if you want learning over log consumption policy, rather than untransformed consumption policy.




% Stationary distribution parameters and preallocation
N = 250;               % set number of agents in cross section to simulate for stat. distr.
smooth_wind = 200;      % smoothing window size (to evaluate stationary distr. convergence)
crit_distr = 1e-1;      % convergence criteria for stationary distr.
T          = max(10*1/deltax,1500);
maxiter_distr = max(1500,T);   % max number of iterations in search for stationary distr.
welf_sim_size = max(200, maxiter_distr - 2*T);

%%%%% Big simulation
sim_periods = maxiter_distr;
nsim = 500; N;

sim_periods_irf = 1;


yi = exp(normrnd( lbar,sigma_l,sim_periods,nsim)); % preallocate all y-shocks
unemp_shocks = rand(sim_periods, nsim) < p_unemp;
yi(unemp_shocks == 1) = w_unemp*mean(yi(:));

eps_etai = normrnd(0,1, sim_periods, nsim); % preallocate eta shocks
reset_shocks = rand(sim_periods, nsim); % reset-shocks


% Parameters for interest rate search (bisection method)
% bounds for interst rate (be careful with these)
rLower =  0.035; -delta;
rUpper =  btat.^-1 - 1;

r0 = mean([rLower rUpper]); % initial guess

conv_r = 1;
counter_r = 0;
maxiter_r = 100;             % max number of iterations (search for r)
crit_r = 1e-5;              % stopping criteria when abs(r-rnew)<crit_r  (search for r)

amin = b;
%amax = 45;
%na = 1001;
amax = 20;
na = 1001; round((amax - amin)/0.2);

stateGrid_size = 1001;

%%% keep track of interest rate and market clearing in each iteration

param_fp = [bta,r0,abar,W_cc,kappa,deltax,sigma_c,psi,cov_fun,gam,sigma_l, lbar, delta_eta, delta_inh, delta, alpha,b, btat, amin, amax, na, crit_distr];


%%
%%%%% ENDOGENOUS SIGNALS MODEL

%%%  Find the endogenous signals r

[r_br, rt_br, w_br] = endsignals_r(param_fp, util_type, smooth_wind, T, N, yi, eps_etai, reset_shocks, rLower, rUpper, conv_r, counter_r, maxiter_r, crit_r, run_fsolve, stateGrid_size,log_c, p_unemp, w_unemp, ex_signals);


[xs,~,cstar,~,ctilde] = bc_opt_pol_iid_egm(btat, rt_br, b, log(w_br) + lbar, sigma_l, gam, amin, amax, na, util_type, p_unemp, w_unemp);
[xs,xi] = unique(xs);


% get policy function over an equally spaced grid
stateGrid_size = 1001;
stateGrid = linspace( min(xs(:)), max(xs(:)),stateGrid_size);
stateGrid_step = stateGrid(2) - stateGrid(1);

% ctilde is the optimal choice for consumption if the borrowing
% constraint were relaxed only in the current period.
if log_c == 0
    temp = griddedInterpolant(xs,ctilde(xi),'linear','linear');
else
    temp = griddedInterpolant(xs,log(ctilde(xi)),'linear','linear');
end

mu0_fun = temp;
ctilde = temp(stateGrid);

temp = griddedInterpolant(xs,cstar(xi),'linear','linear');

cstar_fixedpoints = temp(stateGrid);
temp = griddedInterpolant(stateGrid,cstar_fixedpoints,'linear','linear');
%cstar = @(x) temp(x);
cstar = @(x) fast_interp( x, cstar_fixedpoints, stateGrid_step, stateGrid,stateGrid_size);
cstar_brsim = mu0_fun;

rw_pol = (rt_br/(1+rt_br))*stateGrid' + 1/(1+rt_br)*w_br;
rw_pol_interp = griddedInterpolant(stateGrid, rw_pol,'linear','linear');


param_br_sim = [bta,rt_br,abar,W_cc,kappa,deltax,sigma_c,psi,cov_fun,gam,sigma_l, lbar, delta_eta, delta_inh, factor_mpc1, factor_mpc2,0];

tic
[a_sims_br,chat_sims_br, cstar_sims_br, aphat_sims_br,astar_sims_br,eta_sims,sigma_etasq_sims,alpha_star_sims,mpc_br,mpc_br2, mpc_br3, mpc_br4, age_sims_br,prior_sims, chat_prime_sims] = bc_endsignals_sim(w_br,param_br_sim,yi(1:sim_periods,1:nsim),eps_etai(1:sim_periods,1:nsim),reset_shocks(1:sim_periods,1:nsim),b,mu0_fun, ctilde, stateGrid,stateGrid_step,stateGrid_size,cstar,log_c, abar*ones(1,nsim), ex_signals);
toc


%%

yi_infoshock = yi; % preallocate all y-shocks
eps_etai_infoshock = eps_etai; % preallocate eta shocks
reset_shocks_infoshock = reset_shocks; % reset-shocks

reset_shocks_infoshock(end-sim_periods_irf+1,:) = 0;

delta_eta_infoshock = psi_tau; %%% 0.5 replicates the fall in expenditure shares reported by krueger et. al. (2016)
param_br_sim = [bta,rt_br,abar,W_cc,kappa,deltax,sigma_c,psi,cov_fun,gam,sigma_l, lbar, delta_eta, delta_inh, factor_mpc1, factor_mpc2,delta_eta_infoshock];


tic
[a_sims_infoshock_br,chat_sims_infoshock_br, cstar_sims_infoshock_br, aphat_sims_infoshock_br,apstar_infoshock_sims_br,eta_sims_infoshock,sigma_etasq_sims_infoshock,alpha_star_sims_infoshock,mpc_br_infoshock,mpc_br2_infoshock, mpc_br3_infoshock, mpc_br4_infoshock, age_sims_infoshock,prior_sims_infoshock] = bc_endsignals_sim(w_br,param_br_sim,yi_infoshock,eps_etai_infoshock,reset_shocks_infoshock,b,mu0_fun, ctilde, stateGrid,stateGrid_step,stateGrid_size,cstar,log_c, abar*ones(1,nsim), ex_signals);
toc



%%
%%%%% RATIONAL EXPECTATIONS MODEL

[r_re, rt_re, w_re] = re_r(param_fp, util_type, N, smooth_wind, T, yi, reset_shocks, rLower, rUpper, conv_r, counter_r, maxiter_r, crit_r, run_fsolve, stateGrid_size, p_unemp, w_unemp);


[xs,apstar, cstar,~,ctilde] = bc_opt_pol_iid_egm(btat, rt_re, b, log(w_re) + lbar, sigma_l, gam, amin, amax, na, util_type, p_unemp, w_unemp);
[xs,xi] = unique(xs);

% get policy function over an equally spaced grid
stateGrid_size = 1001;
stateGrid = linspace( min(xs(:)), max(xs(:)),stateGrid_size);
stateGrid_step = stateGrid(2) - stateGrid(1);

% ctilde is the optimal choice for consumption if the borrowing
% constraint were relaxed only in the current period.
temp = griddedInterpolant(xs,cstar(xi),'linear','linear');

cstar_fixedpoints = temp(stateGrid);
temp = griddedInterpolant(stateGrid,cstar_fixedpoints,'linear','linear');
%cstar = @(x) temp(x);
cstar = @(x) fast_interp( x, cstar_fixedpoints, stateGrid_step, stateGrid,stateGrid_size);

rw_pol_re = (rt_re/(1+rt_re))*stateGrid' + 1/(1+rt_re)*w_re;
rw_pol_interp_re = griddedInterpolant(stateGrid, rw_pol_re,'linear','linear');


%temp = xs - cstar;

param_re_sim = [bta,rt_re,abar,W_cc,kappa,deltax,sigma_c,psi,cov_fun,gam,sigma_l, lbar, delta_eta, delta_inh, factor_mpc1, factor_mpc2];


tic
[a_sims_re,cstar_sims_re,aphat_sims_re,age_sims_re, mpc_re, mpc_re2, mpc_re3, mpc_re4] = bc_re_sim(w_re,param_re_sim,yi(1:sim_periods,1:nsim),reset_shocks(1:sim_periods,1:nsim),b,cstar,abar*ones(1,nsim));
toc


%%
%%%%% Trembles model

[r_tr, rt_tr, w_tr] = tremble_r(param_fp, util_type, N, smooth_wind, T, yi, reset_shocks, rLower, rUpper, conv_r, counter_r, maxiter_r, crit_r, run_fsolve, stateGrid_size,eps_etai,sigma_tremble, p_unemp, w_unemp);


[xs,apstar, cstar,~,ctilde] = bc_opt_pol_iid_egm(btat, rt_tr, b, log(w_tr) + lbar, sigma_l, gam, amin, amax, na, util_type, p_unemp, w_unemp);
[xs,xi] = unique(xs);

% get policy function over an equally spaced grid
stateGrid_size = 1001;
stateGrid = linspace( min(xs(:)), max(xs(:)),stateGrid_size);
stateGrid_step = stateGrid(2) - stateGrid(1);

% ctilde is the optimal choice for consumption if the borrowing
% constraint were relaxed only in the current period.
temp = griddedInterpolant(xs,cstar(xi),'linear','linear');

cstar_fixedpoints = temp(stateGrid);
temp = griddedInterpolant(stateGrid,cstar_fixedpoints,'linear','linear');
%cstar = @(x) temp(x);
cstar = @(x) fast_interp( x, cstar_fixedpoints, stateGrid_step, stateGrid,stateGrid_size);

%temp = xs - cstar;

param_tr_sim = [bta,rt_tr,abar,W_cc,kappa,deltax,sigma_c,psi,cov_fun,gam,sigma_l, lbar, delta_eta, delta_inh, factor_mpc1, factor_mpc2];


tic
[a_sims_tr,cstar_sims_tr,aphat_sims_tr,age_sims_tr, mpc_tr, mpc_tr2, mpc_tr3, mpc_tr4] = bc_tremble_sim(w_tr,param_tr_sim,yi(1:sim_periods,1:nsim),reset_shocks(1:sim_periods,1:nsim),eps_etai(1:sim_periods,1:nsim),sigma_tremble,b,cstar,abar*ones(1,nsim));
toc



%% Computing moments of interest



K_br = ((r_br + delta)/alpha)^(1/(alpha-1)); % Aggregate demand of capital (from firm k-FOC)
K_re = ((r_re + delta)/alpha)^(1/(alpha-1)); % Aggregate demand of capital (from firm k-FOC)
K_tr = ((r_tr + delta)/alpha)^(1/(alpha-1)); % Aggregate demand of capital (from firm k-FOC)

mu0_fun = cstar_brsim;

return 
compute_moments






